import streamlit as st
import requests
import json
from typing import Dict, Any
import uuid
import sys
from loguru import logger

# 配置loguru日志
logger.remove()  # 移除默认配置

# 控制台日志 - 显示INFO及以上级别
logger.add(
    sys.stdout,
    level="INFO",
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
    colorize=True
)

# 文件日志 - 详细日志，包含DEBUG级别
logger.add(
    "logs/flood_forecast_streamlit_demo.log",
    level="DEBUG",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} | {process.id} | {thread.id} - {message}",
    rotation="10 MB",
    retention="7 days",
    compression="zip",
    encoding="utf-8"
)

# 错误日志单独文件
logger.add(
    "logs/flood_forecast_streamlit_demo_error.log",
    level="ERROR",
    format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} | {process.id} | {thread.id} - {message}\n{exception}",
    rotation="5 MB",
    retention="30 days",
    compression="zip",
    encoding="utf-8"
)

# 页面配置
st.set_page_config(
    page_title="洪水预报参数收集系统",
    page_icon="🌊",
    layout="wide"
)

# API配置
API_BASE_URL = "http://localhost:8001"

def init_session_state():
    """初始化会话状态"""
    if 'session_id' not in st.session_state:
        st.session_state.session_id = str(uuid.uuid4())
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    if 'parameters' not in st.session_state:
        st.session_state.parameters = {}
    if 'collection_status' not in st.session_state:
        st.session_state.collection_status = 'collecting'
    
    # 确保服务器端会话存在
    ensure_server_session()

def ensure_server_session():
    """确保服务器端会话存在"""
    try:
        # 检查会话是否存在
        response = requests.get(
            f"{API_BASE_URL}/flood_forecast/session/{st.session_state.session_id}/status"
        )
        if response.status_code == 404:
            # 会话不存在，通过发送一个空请求来创建会话
            requests.post(
                f"{API_BASE_URL}/flood_forecast/collect_parameters",
                json={
                    "session_id": st.session_state.session_id,
                    "user_input": "初始化会话"
                }
            )
    except requests.exceptions.RequestException:
        # 如果服务器不可用，继续使用本地状态
        pass

def send_user_input(user_input: str) -> Dict[str, Any]:
    """发送用户输入到API"""
    try:
        response = requests.post(
            f"{API_BASE_URL}/flood_forecast/collect_parameters",
            json={
                "session_id": st.session_state.session_id,
                "user_input": user_input
            }
        )
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        st.error(f"API请求失败: {str(e)}")
        return None

def get_final_parameters() -> Dict[str, Any]:
    """获取最终参数JSON"""
    try:
        response = requests.get(
            f"{API_BASE_URL}/flood_forecast/session/{st.session_state.session_id}/final_json"
        )
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        st.error(f"获取最终参数失败: {str(e)}")
        return None

def reset_session():
    """重置会话"""
    try:
        # 先检查会话是否存在
        status_response = requests.get(
            f"{API_BASE_URL}/flood_forecast/session/{st.session_state.session_id}/status"
        )
        
        if status_response.status_code == 404:
            # 会话不存在，直接重置本地状态
            st.session_state.chat_history = []
            st.session_state.parameters = {}
            st.session_state.collection_status = 'collecting'
            st.success("会话已重置")
            return
            
        # 会话存在，调用重置API
        response = requests.post(
            f"{API_BASE_URL}/flood_forecast/session/{st.session_state.session_id}/reset"
        )
        response.raise_for_status()
        st.session_state.chat_history = []
        st.session_state.parameters = {}
        st.session_state.collection_status = 'collecting'
        st.success("会话已重置")
    except requests.exceptions.RequestException as e:
        # 如果重置失败，至少重置本地状态
        st.session_state.chat_history = []
        st.session_state.parameters = {}
        st.session_state.collection_status = 'collecting'
        st.warning(f"服务器重置失败，但本地状态已重置: {str(e)}")

def display_parameters_panel(parameters: Dict[str, Any], missing_params: list):
    """显示参数面板"""
    st.subheader("📋 参数收集状态")
    
    # 显示已收集的参数
    if parameters:
        st.write("**已收集的参数:**")
        for key, value in parameters.items():
            if value is not None:
                st.write(f"✅ {key}: {value}")
    
    # 显示缺失的参数
    if missing_params:
        st.write("**待收集的参数:**")
        for param in missing_params:
            st.write(f"❌ {param}")
    else:
        st.success("🎉 所有参数已收集完成！")

def main():
    """主函数"""
    init_session_state()
    
    # 标题
    st.title("🌊 洪水预报参数收集系统")
    st.markdown("---")
    
    # 布局
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.subheader("💬 对话界面")
        
        # 显示聊天历史
        chat_container = st.container()
        with chat_container:
            for message in st.session_state.chat_history:
                if message['role'] == 'user':
                    st.chat_message("user").write(message['content'])
                else:
                    st.chat_message("assistant").write(message['content'])
        
        # 用户输入
        if st.session_state.collection_status != 'completed':
            user_input = st.chat_input("请输入您的洪水预报需求...")
            
            if user_input:
                # 添加用户消息到历史
                st.session_state.chat_history.append({
                    'role': 'user',
                    'content': user_input
                })
                
                # 发送到API
                with st.spinner("AI正在分析您的需求..."):
                    result = send_user_input(user_input)
                
                if result:
                    # 添加AI回复到历史
                    st.session_state.chat_history.append({
                        'role': 'assistant',
                        'content': result['response']
                    })
                    
                    # 更新状态
                    st.session_state.parameters = result['parameters']
                    st.session_state.collection_status = result['status']
                    
                    # 重新运行以更新界面
                    st.rerun()
        else:
            st.success("✅ 参数收集已完成！")
            
            # 显示最终参数按钮
            if st.button("📄 查看最终参数JSON", type="primary"):
                final_params = get_final_parameters()
                if final_params:
                    st.subheader("最终参数JSON")
                    st.json(final_params['final_parameters'])
                    
                    # 提供下载功能
                    json_str = json.dumps(final_params['final_parameters'], ensure_ascii=False, indent=2)
                    st.download_button(
                        label="💾 下载参数JSON",
                        data=json_str,
                        file_name=f"flood_forecast_params_{st.session_state.session_id[:8]}.json",
                        mime="application/json"
                    )
    
    with col2:
        # 参数状态面板
        if hasattr(st.session_state, 'parameters'):
            missing_params = []
            # 这里可以从API获取缺失参数列表，简化起见直接计算
            required_params = ['reservoir_name', 'forecast_model', 'start_time', 
                             'forecast_duration', 'rainfall_forecast_period', 'target_variable']
            missing_params = [p for p in required_params if not st.session_state.parameters.get(p)]
            
            display_parameters_panel(st.session_state.parameters, missing_params)
        
        # 控制按钮
        st.markdown("---")
        st.subheader("🔧 操作")
        
        if st.button("🔄 重置会话", type="secondary"):
            reset_session()
            st.rerun()
        
        # 会话信息
        st.markdown("---")
        st.subheader("ℹ️ 会话信息")
        st.write(f"**会话ID:** {st.session_state.session_id[:8]}...")
        st.write(f"**状态:** {st.session_state.collection_status}")
        
        # 示例问题
        st.markdown("---")
        st.subheader("💡 示例问题")
        example_questions = [
            "我要做一个大伙房水库的新安江预报，预测未来流量",
            "需要做三峡水库的48小时流量预报",
            "帮我设置一个水位预报，从明天开始"
        ]
        
        for i, question in enumerate(example_questions):
            if st.button(f"示例 {i+1}", key=f"example_{i}"):
                if st.session_state.collection_status != 'completed':
                    # 模拟用户输入
                    st.session_state.chat_history.append({
                        'role': 'user',
                        'content': question
                    })
                    
                    # 发送到API
                    result = send_user_input(question)
                    if result:
                        st.session_state.chat_history.append({
                            'role': 'assistant',
                            'content': result['response']
                        })
                        st.session_state.parameters = result['parameters']
                        st.session_state.collection_status = result['status']
                        st.rerun()

    # 页脚
    st.markdown("---")
    st.markdown(
        """
        <div style='text-align: center; color: #666;'>
            🌊 洪水预报参数收集系统 | 基于LangChain的智能对话式参数收集
        </div>
        """, 
        unsafe_allow_html=True
    )

if __name__ == "__main__":
    main()